from ultralytics import YOLO
import glob
from pathlib import Path
import os
from tqdm import tqdm

# # Load a model
# model = YOLO("/root/yolov8-SAR/runs/obb/train4/weights/best.pt")  # pretrained YOLOv8n model

# # Run batched inference on a list of images
# imgpaths=glob.glob("/root/yolov8-SAR/datasets/SAR/test/JPEGImages/*.jpg")
# imgpaths = sorted(imgpaths, key=lambda x: int(os.path.basename(x).split('.')[0]))

# with open("commit.txt",'w') as f:
#     for imgpath in tqdm(imgpaths):
#         results = model([imgpath])  # return a list of Results objects
#         for result in results:
#             orig_shape=result.orig_shape    # (256,256)
#             img_path=result.path        # abs path
#             for idx,xy4 in enumerate(result.obb.xyxyxyxy):   # 图像中多个bbox
#                 conf=str(result.obb.conf[idx].cpu().detach().numpy())       # conf: tensor([0.5867, 0.3250], device='cuda:0')
#                 xy4=xy4.detach().cpu().numpy().flatten()
#                 xy4=" ".join([str(int(pos)) for pos in xy4])
#                 line=str(Path(img_path).name)+" "+xy4+" "+conf
#                 f.write(line+"\n")

from ultralytics import YOLO
import os
import glob
import cv2
from pathlib import Path

# Load a model
model = YOLO("runs/obb/train4/weights/best.pt")

imgpaths=glob.glob("datasets/SAR/test/JPEGImages/*.jpg")
imgpaths = sorted(imgpaths, key=lambda x: int(os.path.basename(x).split('.')[0]))

with open("commit.txt",'w') as f:
    for imgpath in tqdm(imgpaths):
        img=cv2.imread(imgpath)
        H,W,C=img.shape
        results=model.predict(img, save=False, imgsz=1024, conf=0.5)
        for result in results:
            for idx,xy4 in enumerate(result.obb.xyxyxyxy):   # 图像中多个bbox
                conf=str(result.obb.conf[idx].cpu().detach().numpy())       # conf: tensor([0.5867, 0.3250], device='cuda:0')
                xy4=xy4.detach().cpu().numpy().flatten()
                xy4=" ".join([str(int(pos)) for pos in xy4])
                line=str(Path(imgpath).name)+" "+xy4+" "+conf
                f.write(line+"\n")